from pocketflow import AsyncNode
from scorpio.core.schemas import Message
from scorpio.core.common import AgentConfig
from scorpio.core.prompt import PromptManager
from scorpio.core.mcp import MCPClient

from scorpio.core.common import get_logger
logger = get_logger(__name__)

DEFAULT_SYSTEM_PROMPT="simple.md"
class PrepareAction(AsyncNode):
    async def prep_async(self, shared):
        """Get the search query from the shared store."""
        context = shared.get("context")
        
        question = shared["question"]
        system_prompt= shared.get("system_prompt", DEFAULT_SYSTEM_PROMPT)

        prompt_manager= PromptManager()
        prompt=prompt_manager.get_user_prompt_template(system_prompt)
        if context.get("messages") is None:
            context["messages"]=[]
        messages=context.get("messages")
        system_message=Message(role="system",content=prompt)
        messages.append(system_message)
        user_message=Message(role="user", content=question)
        messages.append(user_message)
        return context, messages
        
    async def exec_async(self, inputs):
        """Search the web for the given query."""
        # Call the search utility function
        context, messages = inputs
        logger.info(f"prepare messages:\n {messages}.")
        agent= context["agent"]
        if agent is None:
            return "failed"
        
        mcp_server = agent["mcp_server"]
        agent_config: AgentConfig = agent["config"]
        # cli: MCPClient= context.get("client")
        cli = MCPClient(
            mcp_server,
            mcp_config=agent_config.mcp_config,
            model_param=agent_config.model_parameter
        )
        await cli.start()
        context["mcp_client"]= cli
        context["llm_client"]= cli.model
        return "prepared"
    
    async def post_async(self, shared, prep_res, exec_res):
        """Save the search results and go back to the decision node."""
        if exec_res != "prepared":
            logger.error(f"failed to prepare agent resources in Pocketflow FlowPreparation") 
        return exec_res